package vip.yeee.memoo.base.swagger.filter;

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.ServletRequest;
import jakarta.servlet.Filter;
import jakarta.servlet.ServletResponse;
import jakarta.servlet.WriteListener;
import jakarta.servlet.annotation.WebFilter;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
import java.util.Base64;

/**
 * description......
 *
 * @author yeeee
 * @since 2025/7/25 17:53
 */
@WebFilter(urlPatterns = {"/v3/api-docs", "/v3/api-docs/*"})
public class ResponseModifyFilter implements Filter {

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest req = (HttpServletRequest) request;
        CachedBodyHttpServletResponse responseWrapper = new CachedBodyHttpServletResponse((HttpServletResponse) response);

        chain.doFilter(request, responseWrapper);

        byte[] body = responseWrapper.getCachedBody();
        String bodyStr = new String(body, response.getCharacterEncoding());

        System.out.println(">>> [Response Log] " + req.getRequestURI() + " 返回内容：");
        System.out.println(bodyStr);

        // 去除前后双引号（如果有）
        String trimmedStr = bodyStr;
        if (trimmedStr.startsWith("\"")) {
            trimmedStr = trimmedStr.substring(1);
        }
        if (trimmedStr.endsWith("\"")) {
            trimmedStr = trimmedStr.substring(0, trimmedStr.length() - 1);
        }

        byte[] outputBytes;
        boolean isBase64 = isBase64(trimmedStr);

        if (isBase64) {
            // Base64合法则解码
            outputBytes = Base64.getDecoder().decode(trimmedStr);
            String decodedStr = new String(outputBytes, response.getCharacterEncoding());
            System.out.println(">>> 解码后的内容：");
            System.out.println(decodedStr);
        } else {
            // 非Base64直接返回原始body
            outputBytes = body;
            System.out.println(">>> 内容非Base64，直接返回原始响应");
        }

        // 写回响应
        response.setContentLength(outputBytes.length);
        ServletOutputStream out = response.getOutputStream();
        out.write(outputBytes);
        out.flush();
    }

    /**
     * 判断字符串是否为合法Base64
     */
    private boolean isBase64(String str) {
        if (str == null || str.isEmpty()) {
            return false;
        }
        // Base64字符串长度必须是4的倍数
        if (str.length() % 4 != 0) {
            return false;
        }
        // 只允许Base64字符，含+ / =号
        if (!str.matches("^[A-Za-z0-9+/=\\r\\n]+$")) {
            return false;
        }
        try {
            Base64.getDecoder().decode(str);
            return true;
        } catch (IllegalArgumentException e) {
            return false;
        }
    }


    /**
     * 这个类包装 HttpServletResponse，缓存响应体，方便我们读和修改
     */
    private static class CachedBodyHttpServletResponse extends HttpServletResponseWrapper {
        private ByteArrayOutputStream cachedContent = new ByteArrayOutputStream();
        private ServletOutputStream outputStream;
        private PrintWriter writer;

        public CachedBodyHttpServletResponse(HttpServletResponse response) {
            super(response);
        }

        @Override
        public ServletOutputStream getOutputStream() {
            if (this.outputStream == null) {
                this.outputStream = new CachedServletOutputStream(cachedContent);
            }
            return this.outputStream;
        }

        @Override
        public PrintWriter getWriter() throws UnsupportedEncodingException {
            if (this.writer == null) {
                this.writer = new PrintWriter(new OutputStreamWriter(cachedContent, getCharacterEncoding()));
            }
            return this.writer;
        }

        public byte[] getCachedBody() {
            // 确保writer缓冲写出
            if (writer != null) {
                writer.flush();
            }
            return cachedContent.toByteArray();
        }

        private static class CachedServletOutputStream extends ServletOutputStream {
            private final OutputStream out;

            public CachedServletOutputStream(OutputStream out) {
                this.out = out;
            }

            @Override
            public void write(int b) throws IOException {
                out.write(b);
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener listener) {
                // no-op
            }
        }
    }
}

